Btw, key can also be a key for e.g. indexing into list or a dict or a list of such keys. For example,

In [8]: toolz.countby('a', [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
Out[8]: {1: 1, 10: 1}

In [9]: toolz.countby('b', [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
Out[9]: {2: 2}

In [10]: toolz.countby(['a', 'b'], [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
Out[10]: {(1, 2): 1, (10, 2): 1}

In [11]: toolz.countby(0, [[1, 2], [10, 2]])
Out[11]: {1: 1, 10: 1}

In [12]: toolz.countby(1, [[1, 2], [10, 2]])
Out[12]: {2: 2}

In [13]: toolz.countby([0, 1], [[1, 2], [10, 2]])
Out[13]: {(1, 2): 1, (10, 2): 1}

How should we declare the type of key? Something like Union[KT, List[KT], Callable[[A], B]]? Do you like KT (following how typing types mappings) or K?

This particular pattern comes up a lot, so perhaps we should make an object to use instead of the Union expression. For example, this could show up in the function definition as def countby(key: Key, .... Key isn't very descriptive. Anything better?

Then, in this case, seq can be different things too.

It's like we have multiple versions of the same countby, as the type of key seems to be linked to the type of seq.

So, we have the following versions if I understand well:

# Case: countby(len, ['cat', 'mouse', 'dog'])
def countby(key: Callable[[A], B], seq: Iterable[A]) -> Dict[B, int]:

# Case: countby('a', [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
def countby(key: K, seq: Iterable[Dict[K, B]]) -> Dict[B, int]:

# Case: countby(['a', 'b'], [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
def countby(key: List[K], seq: Iterable[Dict[K, B]]) -> Dict[Tuple[B,...], int]:

# Case: countby(0, [[1, 2], [10, 2]])
def countby(key: int, seq: Iterable[Iterable[A]]) -> Dict[A, int]:

# Case: countby([0, 1], [[1, 2], [10, 2]])
def countby(key: List[int], seq: Iterable[Iterable[A]]) -> Dict[Tuple[A, ...], int]:

So, in this case, the types aren't so useful for such a generic function, because they would loose that dependency between inputs and output.

What do you think?

I do agree on reusing an object instead of duplicating the Union everywhere though.

In the end, I think the most generic way I could type this is:

def countby(key: KeyLike, seq: Iterable) -> Dict[Any, int]:

Python's type system does not allow making the dependency possible.

Also, perhaps I should use Sequence instead of Iterable. What do you think?

You can have dependencies if you use overload. Something like,

# Case: countby(len, ['cat', 'mouse', 'dog'])
@overload
def countby(key: Callable[[A], B], seq: Iterable[A]) -> Dict[B, int]: ...

# Case: countby('a', [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
@overload
def countby(key: K, seq: Iterable[Dict[K, B]]) -> Dict[B, int]: ...

def countby(key: KeyLike, seq: Iterable) -> Dict[Any, int]: 
   # real definition

The mypy documentation discusses overload here.

Thanks, I've updated that.

only checking the file I modified added type hints to. But it depends on other files: this explains why I had to ignore some errors / make mypy happy elsewhere.